#include "phv_filter.h"

#include <eigen3/Eigen/Core>
#include <eigen3/Eigen/Geometry>

#include "../global.h"

namespace fusion_perception {

PHVFilter::PHVFilter()
{
    init();
}

void PHVFilter::init(void)
{
    auto X = Eigen::Matrix<double, ds_, 1>::Zero();
    auto P = Eigen::Matrix<double, ds_, ds_>::Identity();
    auto F = Eigen::Matrix<double, ds_, ds_>::Identity();
    auto H = Eigen::Matrix<double, dm_, ds_>::Identity();

    kf_.SetStateEstimate(X, P);
    kf_.SetTransitionMatrix(F);
    kf_.SetObservationMatrix(H);
}

void PHVFilter::predict(double dtS)
{
    auto& global = Global::getInstance();

    auto Q = Eigen::Matrix<double, ds_, ds_>::Identity() * global.Q;
    auto R = Eigen::Matrix<double, dm_, dm_>::Identity() * global.R;

    kf_.SetTransitionNoise(Q);
    kf_.SetObservationNoise(R);

    auto F = kf_.GetTransitionMatrix();
    F << 1, dtS, 0, 0, 0, // x
        0, 1, 0, 0, 0, // vx
        0, 0, 1, dtS, 0, // y
        0, 0, 0, 1, 0, // vy
        0, 0, 0, 0, 1; // heading

    kf_.SetTransitionMatrix(F);
    kf_.Predict();
}

PHVFilter::CorrectedPosition PHVFilter::correct(double dtS, PHVFilter::MeasuredPosition measuredPosition)
{
    if (xInitFlags_ == false) {
        xInitFlags_ = true;
        auto X = kf_.GetStateEstimate();
        X << measuredPosition.x, measuredPosition.vx, measuredPosition.y, measuredPosition.vy, measuredPosition.heading;
        kf_.SetStateEstimate(X, kf_.GetStateCovariance());
    }

    predict(dtS);
    return correct(measuredPosition);
}

PHVFilter::CorrectedPosition PHVFilter::correct(PHVFilter::MeasuredPosition measuredPosition)
{
    Eigen::Matrix<double, dm_, 1> Z;
    Z.setZero();
    Z << measuredPosition.x, measuredPosition.vx, measuredPosition.y, measuredPosition.vy, measuredPosition.heading;

    kf_.Correct(Z);
    auto ret = kf_.GetStateEstimate();

    CorrectedPosition correctedPosition;
    correctedPosition.x = ret[0];
    correctedPosition.vx = ret[1];
    correctedPosition.y = ret[2];
    correctedPosition.vy = ret[3];
    correctedPosition.heading = ret[4];

    return correctedPosition;
}

}
